# -*- coding:utf-8 -*-#
# @Time:2023/7/6 17:02
# @Author:Adong
# @Software:PyCharm


"""
LSTM_GAN生成故障声纹demo
"""

import torch
from torch import optim
from torch import nn
from net import LSTM_GAN
import numpy as np
from PIL import Image
import librosa
import random
import matplotlib.pyplot as plt
from torch.utils.data import TensorDataset, DataLoader

# # train_dataset = dataset.makeDataset('train', 'wavV2_to_gray')
# # dataloader = data.DataLoader(dataset=train_dataset, batch_size=8, shuffle=True)
# filepath = r'./data/wav_data_V5/jbfd/局部放电_byq_jbfd_01.wav'
# y, sr = librosa.load(filepath)
# real_data = []  # 真实样本
# for x in random.sample(range(0, round(len(y) - 4096)), 100):
#     piece = y[x:x + 4096]
#     real_data.append(piece)
#
# Train_DS = TensorDataset(torch.Tensor(np.array(real_data)), torch.ones_like(torch.Tensor(np.array(real_data))))
# Train_DL = DataLoader(Train_DS, shuffle=True, batch_size=8)
# device = "cpu"
# # 实例化判别器与生成器
# # z_dim = 100
# # realdata_dim = 64*64
# fakedata_dim = 4096
# LSTMGAN = LSTM_GAN([4096, 1024, 2]).to(device)
#
# # 定义判别器与生成器所使用的优化算法
# optim_disc = optim.Adam(LSTMGAN.discriminator.parameters(), lr=0.001)        #RMSprop,SGD
# optim_gen = optim.Adam(LSTMGAN.generator.parameters(), lr=0.001)
#
# # 定义损失函数
# BCE = nn.BCELoss(reduction="mean")  # 默认生成均值，也可以特定标明生成均值
# MSE = nn.MSELoss(reduction="mean")
#
# real2real_trend = []
# fake2fake_trend = []
# new_fake2fake_trend = []
# disloss = []
# genloss = []
# num_epoch = 100
#
# for epoch in range(num_epoch):
#     # for batch_idx, (x, _) in enumerate(dataloader):
#     print("epoch", epoch)
#     for x, y in Train_DL:
#         real_in = torch.Tensor(x).to(device)        # 加载真样本
#         batch_size = real_in.shape[0]               # 获取批数
#         noise = torch.Tensor(np.random.normal(0, 1, [batch_size, 4096])).to(device)  # 生成噪声
#         fake_in = LSTMGAN.gen_forward(noise).to(device)  # 生成假样本
#
#         real_out = torch.squeeze(LSTMGAN.disc_forward(real_in))  # 判别器对真样本的预测概率
#         fake_out = torch.squeeze(LSTMGAN.disc_forward(fake_in.detach()))  # 需要使用detach来阻止gz进入D的计算图，判别器对假样本的预测概率
#
#         # real_loss = scipy.stats.wasserstein_distance(real_out.cpu().detach().numpy(),torch.ones_like(real_out).cpu().detach().numpy())
#         # fake_loss = scipy.stats.wasserstein_distance(fake_out,torch.zeros_like(fake_out))
#
#         real_loss = BCE(real_out, torch.ones_like(real_out))   # 把真判为真的批次样本平均误差
#         fake_loss = BCE(fake_out, torch.zeros_like(fake_out))   # 把假判为假的批次样本平均误差
#         disc_loss = real_loss + fake_loss  # 判别器判别真假样本的误差之和
#         disloss.append(disc_loss.item())
#
#         disc_loss.backward()                        # 计算disc_loss的梯度
#         optim_disc.step()                  # 更新判别器上的权重，目标是让disc_loss最小
#         LSTMGAN.discriminator.zero_grad()  # 清零判别器的梯度
#
#         real2real = real_out.mean().item()          # 判别器把真判为真的平均准确率
#         real2real_trend.append(real2real)
#         fake2fake = fake_out.mean().item()  # 判别器把假的判成假的的准确率
#         fake2fake_trend.append(fake2fake)
#
#
#         real_out2 = LSTMGAN.disc_forward(real_in)   # 注意，由于在此时判别器上的权重已经被更新过了，所以real_out和fake_out的值会变化，需要重新生成
#         fake_out2 = LSTMGAN.disc_forward(fake_in)
#
#         fake_loss2 = BCE(fake_out2, torch.ones_like(fake_out2))   # 把假判为真的批次样本平均误差
#         gap = MSE(fake_in,real_in)
#         disc_loss2 = fake_loss2 + 100*gap
#         genloss.append(disc_loss2.item())
#
#         disc_loss2.backward()  # 计算disc_loss2的梯度
#         optim_gen.step()                # 更新生成器上的权重，目标是让disc_loss2最小
#         LSTMGAN.generator.zero_grad()  # 清零生成器更新后梯度
#
#         new_fake2fake = fake_out2.mean().item()  # 判别器调整权重后判别假样本的准确率
#         new_fake2fake_trend.append(new_fake2fake)
#
#         print(f'Epoch[{epoch}/{num_epoch}],disc_loss:{disc_loss.data.item()},gen_loss:{disc_loss2.data.item()} '
#               f'real2real: {real_out2.data.mean()},fake2fake: {fake_out2.data.mean()},gap: {gap}')
#     if epoch % 10 == 0:
#         fig = plt.figure()
#         ax1 = plt.subplot(121)
#         ax2 = plt.subplot(122)
#         # 显示每个epoch的loss变化
#         ax1.plot(range(len(real2real_trend)), real2real_trend, label='real2real_trend')
#         # plt.plot(range(len(fake2fake_trend)), fake2fake_trend, color='red', label='fake2fake_trend')
#         ax1.plot(range(len(new_fake2fake_trend)), new_fake2fake_trend, label='new_fake2fake_trend')
#         ax1.legend(loc='upper right')
#         ax1.set_xlabel("epoch")  # x轴上的名字
#         ax1.set_ylabel("acc")  # y轴上的名字
#         ax2.plot(range(len(disloss)), disloss, label='disloss')
#         ax2.plot(range(len(genloss)), genloss, label='genloss')
#         # 简单的设置legend(设置位置)位置在右上角
#         ax2.legend(loc='upper right')
#         ax2.set_xlabel("epoch")  # x轴上的名字
#         ax2.set_ylabel("average loss")  # y轴上的名字
#         # plt.savefig(r'./result_img/' + model_name.split('/')[-1].split('.')[0] + '.png')
#         plt.show()
#
#         '''真假样本画图'''
#         fake_img = fake_in[0].cpu().detach().numpy().reshape(64, 64)
#         real_img = x[0].cpu().detach().numpy().reshape(64, 64)
#         fake_img_norm = (fake_img - np.min(fake_img)) / (np.max(fake_img) - np.min(fake_img)) * 255
#         real_img_norm = (real_img - np.min(real_img)) / (np.max(real_img) - np.min(real_img)) * 255
#         fake_signal = fake_in[0].cpu().detach().numpy()
#         real_signal = x[0].cpu().detach().numpy()
#         fig = plt.figure()
#         ax1 = plt.subplot(221)
#         ax2 = plt.subplot(222)
#         ax3 = plt.subplot(223)
#         ax4 = plt.subplot(224)
#
#         fake = Image.fromarray(fake_img_norm).convert('L')
#         real = Image.fromarray(real_img_norm).convert('L')
#         ax1.imshow(fake, cmap='gray')
#         ax2.imshow(real, cmap='gray')
#         ax3.plot(fake_signal)
#         ax4.plot(real_signal)
#         plt.show()
# # 保存模型
# torch.save(LSTMGAN.generator.state_dict(), 'save_model/generator.pth')
# torch.save(LSTMGAN.discriminator.state_dict(), 'save_model/discriminator.pth')

if __name__ == '__main__':
    import scipy.io.wavfile as _wavefile

    filepath = r'./data/wav_data_V5/jbfd/局部放电_byq_jbfd_01.wav'
    y, sr = librosa.load(filepath)
    LSTMGAN = LSTM_GAN([4096, 1024, 2]).to('cpu')
    LSTMGAN.generator.load_state_dict(torch.load(r'save_model/generator.pth'))
    noise = torch.Tensor(np.random.normal(0, 1, [10, 4096])).to('cpu')  # 生成噪声
    fake_in = LSTMGAN.gen_forward(noise).to('cpu').detach().numpy().squeeze()  # 生成假样本
    fake_in = fake_in.reshape(-1)
    _wavefile.write("generatewav.wav", int(sr), fake_in)
